import pandas as pd

import numpy as np
import math
from eeg_preprocessing import PluginManager,FilterPlugin
import torch

np.warnings.filterwarnings('ignore', category=np.VisibleDeprecationWarning)

# 自定义函数，计算数值的符号。
def sgn(data):
        if data >= 0:
            return 1
        else:
            return 0

    # 计算过零率
def calZeroCrossingRate(data):
        zeroCrossingRate = []
        for i in range(data.shape[0]):
            sum = 0
            for j in range(data.shape[1] - 1):
                sum = sum + np.abs(sgn(data[i][j]) - sgn(data[i][j + 1]))
            zeroCrossingRate.append(float(sum) / (data.shape[1] / 200))  # 200s

        return zeroCrossingRate

def compute_DE(signal):
    variance = np.var(signal, ddof=1)  # 求得方差
    return math.log(2 * math.pi * math.e * variance) / 2  # 微分熵求取公式


def extract_1(raws):
    # 提取EEG数据在五个频段的能量特征
    # delta(0.5-4Hz) theta(4-8Hz) alpha(8-13Hz) beta(13-30Hz) gamma(30-100Hz)
    # 特定频带
    FREQ_BANDS = {"delta": [0.5, 4],
                  "theta": [4, 8],
                  "alpha": [8, 14],
                  "beta": [14, 31],
                  "gamma": [31, 50]}
    # 特征矩阵
    feature_matrix = []
    # 遍历每个raw
    raws.load_data()  # 提取特征需要
    # 生成频谱特征向量
    feature_vector = []
    # 遍历每个频段
    for band in FREQ_BANDS:
        # 提取每个频段的数据，不打印信息
        raw_band = raws.copy().filter(l_freq=FREQ_BANDS[band][0], h_freq=FREQ_BANDS[band][1], verbose=False)
        # 计算能量
        power = np.sum(raw_band.get_data() ** 2, axis=1) / raw_band.n_times
        # 添加到特征向量
        feature_vector.extend(power)
        # 添加到特征矩阵
    feature_matrix.append(feature_vector)
    # 返回特征矩阵
    print("频谱特征矩阵的shape为：{}".format(np.array(feature_matrix).shape))
    # print("频谱特征矩阵内容为：{}".format(np.array(feature_matrix)))
    return np.array(feature_matrix)


def extract_spy(raw):
    # 导联划分
    picks_front = ['FP1', 'FPZ', 'FP2', 'AF3', 'AF4']
    pick_midfront = ['F7', 'F5', 'F3', 'F1', 'FZ', 'F2', 'F4', 'F6', 'F8']
    picks_top = ['FC1', 'FCZ', 'FC2', 'C1', 'CZ', 'C2', 'CP1', 'CPZ', 'CP2']
    picks_left = ['FT7', 'FC5', 'FC3', 'T7', 'C5', 'C3', 'TP7', 'CP5', 'CP3']
    picks_right = ['FC4', 'FC6', 'FT8', 'C4', 'C6', 'T8', 'CP4', 'CP6', 'TP8', ]
    pick_midback = ['P7', 'P5', 'P3', 'P1', 'PZ', 'P2', 'P4', 'P6', 'P8']
    picks_back = ['PO7', 'PO5', 'PO3', 'POZ', 'PO4', 'PO6', 'PO8', 'CB1', 'O1', 'OZ', 'O2', 'CB2']
    picks_list = list([picks_front, pick_midfront, picks_top, picks_left, picks_right, pick_midback, picks_back])
        # 提取EEG数据在五个频段的能量特征
        # delta(0.5-4Hz) theta(4-8Hz) alpha(8-14Hz) beta(14-31Hz) gamma(31-50Hz)
        # 特定频带
    FREQ_BANDS = {"delta": [0.5, 4],
                      "theta": [4, 8],
                      "alpha": [8, 14],
                      "beta": [14, 31],
                      "gamma": [31, 50]}
        # 特征矩阵
    feature_matrix = []
        # 遍历每个raw
    # raw.load_data()  # 提取特征需要

            # 生成频谱特征向量
    feature_vector = []
            # 遍历每个频段
    for band in FREQ_BANDS:
            # 提取每个频段的数据，不打印信息
        raw_band = raw.copy().filter(l_freq=FREQ_BANDS[band][0], h_freq=FREQ_BANDS[band][1], verbose=False)
        # 计算微分熵DE
        data = []
        power = []
        var = []
        p_p = []
        zeroCrossingRate = []
        for pick in picks_list:
            raw_pick = raw_band.get_data(picks=pick)
                    #     print(raw_pick.shape)
                    #             for i in range(62):
                    #  print(raw_pick.shape[0])
            DE_sum = 0
            for i in range(raw_pick.shape[0]):
                DE_sum = DE_sum + compute_DE(raw_pick[i])
            #导入每组数据
            data.append(DE_sum / raw_pick.shape[0])
            power.append(np.mean(np.sum(raw_pick ** 2, axis=1) / raw_band.n_times))
            var.append(np.mean(raw_pick.var(axis=1)))
            p_p.append(np.mean(raw_pick.max(axis=1) - raw_pick.min(axis=1)))
            # zeroCrossingRate.append(np.mean(calZeroCrossingRate(raw_pick)))
        # feature_vector.extend(zeroCrossingRate)
        feature_vector.extend(p_p)
        feature_vector.extend(data)
        feature_vector.extend(var)
        feature_vector.extend(power)
            # 添加到特征矩阵
    feature_matrix.append(feature_vector)
        # 返回特征矩阵
        # print("频谱特征矩阵内容为：{}".format(np.array(feature_matrix)))
    return feature_matrix


def extract_DE(raw):
    # 提取EEG数据在五个频段的能量特征
    # delta(0.5-4Hz) theta(4-8Hz) alpha(8-13Hz) beta(13-30Hz) gamma(30-100Hz)
    # 特定频带
    FREQ_BANDS = {"delta": [0.5, 4],
                  "theta": [4, 8],
                  "alpha": [8, 14],
                  "beta": [14, 31],
                  "gamma": [31, 50]}
    # 特征矩阵
    feature_matrix = []
    # 生成频谱特征向量
    feature_vector = []
    # 遍历每个频段
    for band in FREQ_BANDS:
        # 提取每个频段的数据，不打印信息
        raw_band = raw.copy().filter(l_freq=FREQ_BANDS[band][0], h_freq=FREQ_BANDS[band][1], verbose=False)
        # 计算微分熵DE
        band_data = []
        data=raw_band.get_data()
        for i in range(62):
            DE = compute_DE(data[i])
            band_data.append(DE)
        feature_vector.extend(band_data)
    feature_matrix.append(feature_vector)
    # 返回特征矩阵
    return feature_matrix




def extract_spy5(raw):
    # 导联划分
    picks_front = ['FP1', 'FPZ', 'FP2', 'AF3', 'AF4']
    pick_midfront = ['F7', 'F5', 'F3', 'F1', 'FZ', 'F2', 'F4', 'F6', 'F8']
    picks_top = ['FC1', 'FCZ', 'FC2', 'C1', 'CZ', 'C2', 'CP1', 'CPZ', 'CP2']
    picks_left = ['FT7', 'FC5', 'FC3', 'T7', 'C5', 'C3', 'TP7', 'CP5', 'CP3']
    picks_right = ['FC4', 'FC6', 'FT8', 'C4', 'C6', 'T8', 'CP4', 'CP6', 'TP8', ]
    pick_midback = ['P7', 'P5', 'P3', 'P1', 'PZ', 'P2', 'P4', 'P6', 'P8']
    picks_back = ['PO7', 'PO5', 'PO3', 'POZ', 'PO4', 'PO6', 'PO8', 'CB1', 'O1', 'OZ', 'O2', 'CB2']
    picks_list = list([picks_front, pick_midfront, picks_top, picks_left, picks_right, pick_midback, picks_back])
        # 提取EEG数据在五个频段的能量特征
        # delta(0.5-4Hz) theta(4-8Hz) alpha(8-14Hz) beta(14-31Hz) gamma(31-50Hz)
        # 特定频带
    FREQ_BANDS = {"delta": [0.5, 4],
                      "theta": [4, 8],
                      "alpha": [8, 14],
                      "beta": [14, 31],
                      "gamma": [31, 50]}
    # 特征矩阵
    feature_matrix = []
    # 生成频谱特征向量
    feature_vector = []
            # 遍历每个频段
    for band in FREQ_BANDS:
                # 提取每个频段的数据，不打印信息
        raw_band = raw.copy().filter(l_freq=FREQ_BANDS[band][0], h_freq=FREQ_BANDS[band][1], verbose=False)
        # 计算微分熵DE
        data = []
        power = []
        var = []
        p_p = []
        zeroCrossingRate = []
        for pick in picks_list:#遍历脑区
            raw_pick = raw_band.get_data(picks=pick)
            DE_sum = 0   #计算DE
            for i in range(raw_pick.shape[0]):
                DE_sum = DE_sum + compute_DE(raw_pick[i])
            data.append(DE_sum / raw_pick.shape[0])
            power.append(np.mean(np.sum(raw_pick ** 2, axis=1) /
                                 raw_band.n_times))
            var.append(np.mean(raw_pick.var(axis=1)))
            p_p.append(np.mean(raw_pick.max(axis=1) - raw_pick.min(axis=1)))
            zeroCrossingRate.append(np.mean(calZeroCrossingRate(raw_pick)))
        feature_vector.extend(zeroCrossingRate)
        feature_vector.extend(p_p)
        feature_vector.extend(data)
        feature_vector.extend(var)
        feature_vector.extend(power)
    # 添加到特征矩阵
    feature_matrix.append(feature_vector)
    # 返回特征矩阵
    return feature_matrix